package org.xi.maple.sink

import org.apache.commons.lang3.StringUtils
import org.apache.commons.text.StringSubstitutor
import org.apache.spark.sql.{Dataset, Row, SparkSession}
import org.slf4j.{Logger, LoggerFactory}
import org.xi.maple.api.MapleSink

import scala.collection.JavaConverters._
import scala.collection.mutable

class File extends MapleSink[FileConfig] {

  private val log: Logger = LoggerFactory.getLogger(classOf[File])

  val allowedURISchema = List("file://", "hdfs://")
  val defaultUriSchema = "hdfs://"

  override def checkConfig: (Boolean, String) = {
    if (StringUtils.isBlank(config.getPath)) {
      (false, "please specify [path] as non-empty string")
    } else if (!config.getPath.startsWith("/") && !allowedURISchema.exists(schema => config.getPath.startsWith(schema))) {
      (false, s"invalid path URI, please set the following allowed schemas: ${allowedURISchema.mkString(", ")}")
    } else {
      (true, "")
    }
  }

  def output(spark: SparkSession, argsMap: mutable.Map[String, String], ds: Dataset[Row]): Unit = {
    val writer = ds.write.mode(config.getSaveMode)

    if (config.getPartitionBy != null && !config.getPartitionBy.isEmpty) {
      val partitionKeys = config.getPartitionBy.asScala
      writer.partitionBy(partitionKeys: _*)
    }

    if (config.getOptions != null && !config.getOptions.isEmpty) {
      writer.options(config.getOptions)
    }
    val substitutor = new StringSubstitutor(argsMap.asJava)
    val path = if (config.getPath.startsWith("/")) {
      substitutor.replace(defaultUriSchema + config.getPath)
    } else {
      substitutor.replace(config.getPath)
    }
    log.info(s"Save data to file, path: $path")

    config.getSerializer match {
      case "csv" => writer.csv(path)
      case "json" => writer.json(path)
      case "parquet" => writer.parquet(path)
      case "text" => writer.text(path)
      case "orc" => writer.orc(path)
      case _ => writer.format(config.getSerializer).save(path)
    }
  }
}
